# encoding: utf8
from websocket import create_connection # pip install websocket-client
from time import sleep
import json

from traceroute import do_traceroute
from ping import do_ping
from curl import do_curl


def execute_task(msg_dict):
    """分析收到的消息，执行任务

    :param str msg_dict:
        格式：{"id": "task id", "optype": "ping", "url": "www.baidu.com"}

    :return: json or dict
    """
    print('execute_task: %s' % msg_dict)
    msg = json.loads(msg_dict)
    print(type(msg), msg)
    if msg['optype'] == 'traceroute':
        result = do_traceroute(msg)
    elif msg['optype'] == 'ping':
        result = do_ping(msg)
    elif msg['optype'] == 'curl':
        result = do_curl(msg)
    else: 
        result = msg
    result['probe'] = 'probe_1'
    return result


ws = create_connection("ws://localhost:8080/websocket")
while True:
    print("Waite for task...")
    msg = ws.recv()
    print('receive message: %s' % msg)
    result = execute_task(msg)
    if type(result) is not str:
        result = json.dumps(result)
    ws.send(result)
    print('Send back result: %s' % result)

